import sys
import akshare as ak
import pandas as pd
from datetime import datetime

import akshare as ak
import pandas as pd


def export_mg_weekly(code, year, output_file):
    """
    获取标普500指数周收盘数据并写入CSV
    :param year: 目标年份（如2024）
    :param output_file: 输出文件名（如"sp500_weekly_2024.csv"）
    """
    try:
        # 获取标普500日线数据（接口代码为.INX）
        df = ak.index_us_stock_sina(symbol=code)  # 网页[8]确认标普500代码

         # 统一日期格式和类型
        df["date"] = pd.to_datetime(df["date"])  # 强制转换为 datetime64[ns][6](@ref)
        start_date = pd.to_datetime(f"{int(year)-1}-12-31")
        end_date = pd.to_datetime(f"{year}-12-31")
        
        # 按日期范围过滤
        df = df[(df["date"] >= start_date) & (df["date"] <= end_date)]
        
        # 生成周线数据（取每周最后一个交易日）
        df.set_index("date", inplace=True)
        first_day = df["close"].iloc[[0]].reset_index()

        friday_data = df["close"].resample("W-FRI").last().reset_index()

        weekly_df = pd.concat([first_day, friday_data])
        weekly_df.columns = ["日期", "点数"]
        
        # 保存到 CSV
        weekly_df.to_csv(output_file, index=False, encoding="utf-8-sig")
        print(f"数据已保存至 {output_file}")
    except Exception as e:
        print(f"数据获取失败：{str(e)}")


def export_hs_weekly_data(code, year, output_file):
    try:
        # 获取恒生指数日线数据
        df = ak.stock_hk_index_daily_sina(symbol=code)
        
        # 统一日期格式和类型
        df["date"] = pd.to_datetime(df["date"])  # 强制转换为 datetime64[ns][6](@ref)
        start_date = pd.to_datetime(f"{int(year)-1}-12-31")
        end_date = pd.to_datetime(f"{year}-12-31")
        
        # 按日期范围过滤
        df = df[(df["date"] >= start_date) & (df["date"] <= end_date)]
        
        # 生成周线数据（取每周最后一个交易日）
        df.set_index("date", inplace=True)

        first_day = df["close"].iloc[[0]].reset_index()

        friday_data = df["close"].resample("W-FRI").last().reset_index()

        weekly_df = pd.concat([first_day, friday_data])
        weekly_df.columns = ["日期", "点数"]
        
        # 保存到 CSV
        weekly_df.to_csv(output_file, index=False, encoding="utf-8-sig")
        print(f"数据已保存至 {output_file}")

    except Exception as e:
        print(f"数据获取失败：{str(e)}")


def export_index_weekly(index_code, year, output_file):
    """
    获取指定指数的每周收盘数据并保存为CSV文件
    :param index_code: 指数代码（示例：上证指数"000001"，沪深300"000300"）
    :param year: 目标年份（如2024）
    :param output_file: 输出文件名（如"weekly_index.csv"）
    """
    try:
        # 设置时间范围
        start_date = f"{int(year)-1}1231"
        end_date = f"{year}1231"


        print(start_date)
        # 获取日线数据[1,6](@ref)
        df = ak.index_zh_a_hist(
            symbol=index_code,
            period="daily",
            start_date=start_date,
            end_date=end_date
        )
        
        # 数据预处理[4](@ref)
        df['日期'] = pd.to_datetime(df['日期'])
        df.set_index('日期', inplace=True)

        first_day = df.iloc[[0]][['收盘']].rename(columns={'收盘': '点数'})
        
        # 生成周线数据（取每周最后一个交易日）[6](@ref)
        weekly_df = df.resample('W-FRI').last()  # 以周五为周结束日
        friday_data = weekly_df[['收盘']].rename(columns={'收盘': '点数'})

        weekly_data = pd.concat([first_day, friday_data])
        
        # 保存到CSV[6](@ref)
        weekly_data.reset_index().to_csv(output_file, index=False, encoding='utf-8-sig')
        print(f"数据已保存至 {output_file}，共 {len(weekly_data)} 周记录")
        
    except Exception as e:
        print(f"数据获取失败：{str(e)}")


# https://akshare.akfamily.xyz/data/index/index.html
# hs300: 000300
# 中证500: 000905
# 中证1000: 000852
# 中证红利: 000922
# 红利低波100: 930955
# 创业板: 399006
# 科创50: 000688

# 恒生指数: HSI
# 恒生科技: HSTECH
# 中概互联: H30533

# 标普500: .INX
# 纳指100: .NDX

code=sys.argv[1] # "000001"
year=sys.argv[2] # 2025
output_file=f"weekly_{code}_{year}.csv"

if __name__ == "__main__":
    if code.startswith("HS"):
        export_hs_weekly_data(
            code=code,
            year=year,
            output_file=output_file,
        )
    elif code.startswith("."):
        export_mg_weekly(
            code=code,
            year=year,
            output_file=output_file,
        )
    else:
        export_index_weekly(
            index_code=code,
            year=year,
            output_file=output_file,
        )

